/*
  Simplified BSD License (BSD)

  Copyright (c) 2014, Attila Suszter
  All rights reserved.

  Redistribution and use in source and binary forms, with or without modification, 
  are permitted provided that the following conditions are met:

  * Redistributions of source code must retain the above copyright notice, this 
  list of conditions and the following disclaimer.

  * Redistributions in binary form must reproduce the above copyright notice, this 
  list of conditions and the following disclaimer in the documentation and/or 
  other materials provided with the distribution.

  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 
  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 
  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 
  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR 
  ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 
  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 
  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON 
  ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 
  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 
  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include <iostream>
#include <fstream>
#include <iomanip>
#include <map>
#include "pin.H"
#include "instlib.H"

namespace W {
#include <Windows.h>
}

using std::map;

//
// Global variables
//
static UINT32         ValueToWatch;
static UINT32         WatchFlags;
static UINT32         LogFlags;
static UINT32         ImageFlags;
static UINT32         RelevantInst;
string                Module;
string                DumpPath;
INSTLIB::FILTER       Filter;
wofstream             OutFile;
map<ADDRINT,ADDRINT>  MapOfInstPointers;

#define MAX_DUMPPATH 256

// Watch flags
#define WATCH_REGISTERS     1
#define WATCH_MEMORY_READS  2

// Log flags
#define LOG_STACK_DUMP  1

// Image Flags
#define IMAGE    1
#define NO_IMAGE 2

VOID Dump(ADDRINT address, int size)
{
    for (int i = 0; i < size; i++)
    {
        UINT8 byte;
        size_t size;

        if (i % 16 == 0)
        {
            OutFile << "    " << setw(8) << (ADDRINT)(address+i) << " ";
        }

        size = PIN_SafeCopy(&byte, (void*)(address+i), sizeof(UINT8));

        if (size != sizeof(UINT8))
        {
            OutFile << endl;
            return;
        }

        OutFile << setw(2) << byte << " ";

        if (i % 16 == 15)
        {
            OutFile << endl;
            continue;
        }
    }
    OutFile << endl;
}

VOID BeforeInstruction(ADDRINT reg_eax, ADDRINT reg_ebx, ADDRINT reg_ecx, ADDRINT reg_edx,
                       ADDRINT reg_esi, ADDRINT reg_edi, ADDRINT reg_eip, ADDRINT reg_esp,
                       ADDRINT reg_ebp, const char* disasm, BOOL ignore_match)
{
    bool match = false;

    if (ignore_match)
    {
        match = true;
    }
    else
    {
        if (RelevantInst)
        {
            if ( (strstr(disasm, "eax") != NULL && reg_eax == ValueToWatch) || 
                 (strstr(disasm, "ebx") != NULL && reg_ebx == ValueToWatch) ||
                 (strstr(disasm, "ecx") != NULL && reg_ecx == ValueToWatch) ||
                 (strstr(disasm, "edx") != NULL && reg_edx == ValueToWatch) ||
                 (strstr(disasm, "esi") != NULL && reg_esi == ValueToWatch) ||
                 (strstr(disasm, "edi") != NULL && reg_edi == ValueToWatch) )
            {
                match = true;
            }
        }
        else
        {
            if (reg_eax == ValueToWatch || reg_ebx == ValueToWatch || reg_ecx == ValueToWatch ||
                reg_edx == ValueToWatch || reg_esi == ValueToWatch || reg_edi == ValueToWatch)
            {
                match = true;
            }
        }
    }

    if (match)
    {
        OutFile << "eax=" << setw(8) << reg_eax << " ebx=" << setw(8) << reg_ebx << " ecx=" << setw(8) << reg_ecx << " edx=" << setw(8) << reg_edx << " esi=" << setw(8) << reg_esi << " edi=" << setw(8) << reg_edi << " eip=" << setw(8) << reg_eip << " esp=" << setw(8) << reg_esp << " ebp=" << setw(8) << reg_ebp << endl;

        OutFile << setw(8) << reg_eip << " " << disasm << endl;
        
        if (LogFlags&LOG_STACK_DUMP)
        {
            Dump(reg_esp, 0x48);
        }

        if (DumpPath.empty() == FALSE)
        {
            // Dump memory
            W::MEMORY_BASIC_INFORMATION page;

            if (W::VirtualQuery((W::LPCVOID)reg_eip, &page, sizeof(W::MEMORY_BASIC_INFORMATION)) != 0)
            {
                FILE* dump;
                char fname[MAX_DUMPPATH+16];

                sprintf(fname, "%s\\%08x.dump", DumpPath.c_str(), page.BaseAddress);

                dump = fopen(fname, "wb");

                fwrite(page.BaseAddress, sizeof(char), page.RegionSize, dump);

                fclose(dump);
            }
            else
            {
                OutFile << "Error when dumping memory. EIP=" << reg_eip << " GetLastError=" <<  W::GetLastError() << endl;
            }
        }

        MapOfInstPointers.insert(pair<ADDRINT,ADDRINT>(reg_eip, reg_eip));
    }
}

VOID BeforeReadInstruction(ADDRINT mem_read_ea, ADDRINT reg_eax, ADDRINT reg_ebx, ADDRINT reg_ecx, ADDRINT reg_edx,
                           ADDRINT reg_esi, ADDRINT reg_edi, ADDRINT reg_eip, ADDRINT reg_esp,
                           ADDRINT reg_ebp, const char* disasm)
{
    ADDRINT value;
    size_t size;

    size = PIN_SafeCopy(&value, (void*)mem_read_ea, sizeof(ADDRINT));

    value = _byteswap_ulong(value);

    if (size == sizeof(ADDRINT) && value == ValueToWatch)
    {
        OutFile << "Memory read [" << setw(8) << mem_read_ea << "]=" << value << endl;

        BeforeInstruction(reg_eax, reg_ebx, reg_ecx, reg_edx, reg_esi, reg_edi, reg_eip, reg_esp, reg_ebp, disasm, TRUE);
    }
}

VOID Instruction(INS ins)
{
    // Insert a call to BeforeReadInstruction before any instruction that reads memory
    if ((WatchFlags&WATCH_MEMORY_READS) && INS_IsMemoryRead(ins) && INS_MemoryReadSize(ins) == 4)
    {
        INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)BeforeReadInstruction, IARG_MEMORYREAD_EA,
            IARG_REG_VALUE, REG_EAX, IARG_REG_VALUE, REG_EBX, IARG_REG_VALUE, REG_ECX, IARG_REG_VALUE, REG_EDX,
            IARG_REG_VALUE, REG_ESI, IARG_REG_VALUE, REG_EDI, IARG_REG_VALUE, REG_EIP, IARG_REG_VALUE, REG_ESP,
            IARG_REG_VALUE, REG_EBP, IARG_PTR, (new string(INS_Disassemble(ins)))->c_str(), IARG_END);
    }

    // Insert a call to BeforeInstruction before every instruction
    if (WatchFlags&WATCH_REGISTERS)
    {
        INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)BeforeInstruction,
            IARG_REG_VALUE, REG_EAX, IARG_REG_VALUE, REG_EBX, IARG_REG_VALUE, REG_ECX, IARG_REG_VALUE, REG_EDX,
            IARG_REG_VALUE, REG_ESI, IARG_REG_VALUE, REG_EDI, IARG_REG_VALUE, REG_EIP, IARG_REG_VALUE, REG_ESP,
            IARG_REG_VALUE, REG_EBP, IARG_PTR, (new string(INS_Disassemble(ins)))->c_str(), IARG_BOOL, FALSE, IARG_END);
    }
}

VOID Branch(TRACE trace, VOID *v)
{
    if (!Filter.SelectTrace(trace))
        return;

    for (BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl))
    {
        for (INS ins = BBL_InsHead(bbl); INS_Valid(ins); ins = INS_Next(ins))
        {
            IMG img = IMG_FindByAddress(INS_Address(ins));

            if ((ImageFlags&IMAGE) && IMG_Valid(img))
            {
                string img_name = IMG_Name(img);

                if (Module.empty() == FALSE && img_name.empty() == FALSE &&
                    // Restrict instrumentation to user-supecified module (if any)
                    img_name.find(Module.c_str()) == string::npos)
                {
                    return;
                }

                Instruction(ins);
            }
            else if ((ImageFlags&NO_IMAGE) && (IMG_Valid(img) == false))
            {
                Instruction(ins);
            }
        }
    }
}

//
// This function is called when the application exits
//
VOID Fini(INT32 code, VOID *v)
{
    // Write to a file since cout and cerr maybe closed by the application

    OutFile << endl;
    OutFile << "Summary of Instruction Pointers:" << endl;

    for (map<ADDRINT,ADDRINT>::iterator it = MapOfInstPointers.begin(); it != MapOfInstPointers.end(); it++)
    {
        OutFile << "  " << setw(8) << it->first << endl;
    }

    OutFile.close();
}

//
// Print help
//
INT32 Usage()
{
    cerr << "Trace And Watch for Windows 32-bit" << endl;
    cerr << endl << KNOB_BASE::StringKnobSummary() << endl;
    return -1;
}

// Arguments
//
//   -o <filename>          - specify output log filename, default: TraceAndWatch.log
//   -v <hex value>         - specify value to watch, default: 0x41414141
//   -wf <flags>            - specify the watch flags, default: 1
//   -lf <flags>            - specify the log flags, default: 1
//   -if <flags>            - specify the image flags, default: 1
//   -r <0|1>               - specify if relevant instructions are to be logged only, default 0
//   -m <module>            - specify module to restrict instrumentation to
//   -d <path>              - specify the path of the dump files
//   -filter_no_shared_libs - do not instrument shared libraries
//

KNOB<string> KnobOutputFile(KNOB_MODE_WRITEONCE, "pintool",
    "o", "TraceAndWatch.log", "specify output file name");

KNOB<string> KnobValueToWatch(KNOB_MODE_WRITEONCE, "pintool",
    "v", "41414141", "specify value to watch");

KNOB<UINT32> KnobWatchFlags(KNOB_MODE_WRITEONCE, "pintool",
    "wf", "1", "specify the watch flags");

KNOB<UINT32> KnobLogFlags(KNOB_MODE_WRITEONCE, "pintool",
    "lf", "0", "specify the log flags");

KNOB<UINT32> KnobImageFlags(KNOB_MODE_WRITEONCE, "pintool",
    "if", "1", "specify the image flags");

KNOB<UINT32> KnobRelevant(KNOB_MODE_WRITEONCE, "pintool",
    "r", "0", "specify if relevant instructions are to be logged only");

KNOB<string> KnobModule(KNOB_MODE_WRITEONCE, "pintool",
    "m", "", "specify module to restrict instrumentation to");

KNOB<string> KnobDumpPath(KNOB_MODE_WRITEONCE, "pintool",
    "d", "", "specify the path of the dump files");

int main(int argc, char * argv[])
{
    // Initialize pin
    if (PIN_Init(argc, argv))
    {
        return Usage();
    }

    OutFile.open(KnobOutputFile.Value().c_str());

    ValueToWatch = strtoul(KnobValueToWatch.Value().c_str(), NULL, 16);

    if (ValueToWatch == 0)
    {
        OutFile << "Error when parsing parameter -v" << endl;
        OutFile.close();
        return -1;
    }

    WatchFlags = KnobWatchFlags.Value();
    LogFlags = KnobLogFlags.Value();
    ImageFlags = KnobImageFlags.Value();
    RelevantInst = KnobRelevant.Value();
    Module = KnobModule.Value();
    DumpPath = KnobDumpPath.Value();

    OutFile.fill('0');
    OutFile << hex;

    OutFile << "Trace And Watch Pintool for Windows 32-bit." << endl;
    OutFile << "Value to watch: " << ValueToWatch << endl;
    OutFile << "Watch flags: " << WatchFlags << endl;
    OutFile << "Log flags: " << LogFlags << endl;
    OutFile << "Image flags: " << ImageFlags << endl;
    OutFile << "Relevant Instructions: " << RelevantInst << endl;

    if (Module.empty() == FALSE)
    {
        OutFile << "Module to restrict instrumentation to: " <<  Module.c_str() << endl;
    }
    if (DumpPath.empty() == FALSE)
    {
        if (DumpPath.length() > MAX_DUMPPATH)
        {
            OutFile << "Error: Path of the dump files is too long." << endl;
            OutFile.close();
            return -1;
        }
        OutFile << "Path of the dump files: " <<  DumpPath.c_str() << endl;
    }
    OutFile << endl;

    // Register Instruction to be called to instrument instructions
    TRACE_AddInstrumentFunction(Branch, 0);

    // Register Fini to be called when the application exits
    PIN_AddFiniFunction(Fini, 0);

    Filter.Activate();

    // Start the program, never returns
    PIN_StartProgram();

    return 0;
}
